import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
 
 
ONNX_MODEL = 'best.onnx' # onnx 模型的路径
RKNN_MODEL = './best.rknn'  # 转换后的 RKNN 模型保存路径
DATASET = './pic_path.txt'   # 数据集文件路径
 
QUANTIZE_ON = True   # 是否进行量化
 
if __name__ == '__main__':
 
	# 创建 RKNN 对象
	rknn = RKNN(verbose=True)
 
    # 检查 ONNX 模型文件是否存在
	if not os.path.exists(ONNX_MODEL):
		print('model not exist')
		exit(-1)
 
	# 配置模型预处理参数
	print('--> Config model')
	rknn.config(mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588')
	print('done')
 
	# 加载 ONNX 模型
	print('--> Loading model')
	ret = rknn.load_onnx(model=ONNX_MODEL)
	if ret != 0:
		print('Load yolov5 failed!')
		exit(ret)
	print('done')
 
	# 构建模型
	print('--> Building model')
	ret = rknn.build(do_quantization=QUANTIZE_ON, dataset=DATASET)
	if ret != 0:
		print('Build yolov5 failed!')
		exit(ret)
	print('done')
 
	# 导出 RKNN 模型
	print('--> Export RKNN model')
	ret = rknn.export_rknn(RKNN_MODEL)
	if ret != 0:
		print('Export yolov5rknn failed!')
		exit(ret)
	print('done')
 
